from flask import Blueprint, request, jsonify
from database.config import get_postgres_session
from database.models import Document, DocumentSlice, KnowledgeBase, User
from sqlalchemy.orm import joinedload
import uuid

doc_bp = Blueprint('document', __name__)

@doc_bp.route('/knowledge-bases/<kb_id>/documents', methods=['GET'])
def get_documents_by_kb(kb_id):
    """获取指定知识库的文档列表"""
    session = get_postgres_session()
    try:
        search_term = request.args.get('search', '').lower()
        
        query = session.query(Document).options(
            joinedload(Document.author)
        ).filter_by(knowledge_base_id=kb_id)
        
        if search_term:
            query = query.filter(
                Document.title.ilike(f'%{search_term}%')
            )
        
        documents = query.all()
        
        result = []
        for doc in documents:
            result.append({
                'id': doc.id,
                'title': doc.title,
                'type': doc.type,
                'status': doc.status,
                'lastUpdated': doc.updated_at.strftime('%Y-%m-%d') if doc.updated_at else None,
                'createdAt': doc.created_at.strftime('%Y-%m-%d') if doc.created_at else None,
                'fileType': doc.file_type,
                'fileSize': doc.file_size,
                'pageCount': doc.page_count,
                'author': {
                    'id': doc.author.id,
                    'name': doc.author.name
                } if doc.author else None
            })
        
        return jsonify({
            'success': True,
            'data': result
        }), 200
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500
    finally:
        session.close()

@doc_bp.route('/documents/<doc_id>', methods=['GET'])
def get_document(doc_id):
    """获取文档详情"""
    session = get_postgres_session()
    try:
        doc = session.query(Document).options(
            joinedload(Document.author),
            joinedload(Document.knowledge_base)
        ).filter_by(id=doc_id).first()
        
        if not doc:
            return jsonify({
                'success': False,
                'error': '文档不存在'
            }), 404
        
        result = {
            'id': doc.id,
            'title': doc.title,
            'type': doc.type,
            'status': doc.status,
            'content': doc.content,
            'fileType': doc.file_type,
            'fileSize': doc.file_size,
            'pageCount': doc.page_count,
            'filePath': doc.file_path,
            'lastUpdated': doc.updated_at.strftime('%Y-%m-%d') if doc.updated_at else None,
            'createdAt': doc.created_at.strftime('%Y-%m-%d') if doc.created_at else None,
            'author': {
                'id': doc.author.id,
                'name': doc.author.name
            } if doc.author else None,
            'knowledgeBase': {
                'id': doc.knowledge_base.id,
                'name': doc.knowledge_base.name
            } if doc.knowledge_base else None
        }
        
        return jsonify({
            'success': True,
            'data': result
        }), 200
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500
    finally:
        session.close()

@doc_bp.route('/documents/<doc_id>/slices', methods=['GET'])
def get_document_slices(doc_id):
    """获取文档切片列表"""
    session = get_postgres_session()
    try:
        slices = session.query(DocumentSlice).filter_by(
            document_id=doc_id
        ).order_by(DocumentSlice.slice_index).all()
        
        result = []
        for slice_obj in slices:
            result.append({
                'id': slice_obj.id,
                'question': slice_obj.question,
                'answer': slice_obj.answer,
                'sliceIndex': slice_obj.slice_index,
                'createdAt': slice_obj.created_at.strftime('%Y-%m-%d') if slice_obj.created_at else None
            })
        
        return jsonify({
            'success': True,
            'data': result
        }), 200
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500
    finally:
        session.close()

@doc_bp.route('/knowledge-bases/<kb_id>/documents', methods=['POST'])
def create_document(kb_id):
    """创建新文档"""
    session = get_postgres_session()
    try:
        # 检查知识库是否存在
        kb = session.query(KnowledgeBase).filter_by(id=kb_id).first()
        if not kb:
            return jsonify({
                'success': False,
                'error': '知识库不存在'
            }), 404
        
        data = request.get_json()
        
        if not data or not data.get('title'):
            return jsonify({
                'success': False,
                'error': '文档标题不能为空'
            }), 400
        
        # 生成新的文档ID
        new_id = f"doc{uuid.uuid4().hex[:8]}"
        
        # 创建文档
        doc = Document(
            id=new_id,
            title=data['title'],
            type=data.get('type', 'document'),
            status=data.get('status', 'draft'),
            content=data.get('content', ''),
            file_type=data.get('fileType', 'PDF'),
            file_size=data.get('fileSize', '0 KB'),
            page_count=data.get('pageCount', 1),
            knowledge_base_id=kb_id,
            author_id=data.get('authorId', 1)  # 默认使用管理员用户
        )
        
        session.add(doc)
        session.commit()
        
        return jsonify({
            'success': True,
            'data': {
                'id': doc.id,
                'title': doc.title,
                'type': doc.type,
                'status': doc.status
            }
        }), 201
        
    except Exception as e:
        session.rollback()
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500
    finally:
        session.close()

@doc_bp.route('/documents/<doc_id>', methods=['PUT'])
def update_document(doc_id):
    """更新文档"""
    session = get_postgres_session()
    try:
        doc = session.query(Document).filter_by(id=doc_id).first()
        
        if not doc:
            return jsonify({
                'success': False,
                'error': '文档不存在'
            }), 404
        
        data = request.get_json()
        
        if data.get('title'):
            doc.title = data['title']
        if 'content' in data:
            doc.content = data['content']
        if data.get('type'):
            doc.type = data['type']
        if data.get('status'):
            doc.status = data['status']
        if data.get('fileType'):
            doc.file_type = data['fileType']
        if data.get('fileSize'):
            doc.file_size = data['fileSize']
        if data.get('pageCount'):
            doc.page_count = data['pageCount']
        
        session.commit()
        
        return jsonify({
            'success': True,
            'data': {
                'id': doc.id,
                'title': doc.title,
                'type': doc.type,
                'status': doc.status
            }
        }), 200
        
    except Exception as e:
        session.rollback()
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500
    finally:
        session.close()

@doc_bp.route('/documents/<doc_id>', methods=['DELETE'])
def delete_document(doc_id):
    """删除文档"""
    session = get_postgres_session()
    try:
        doc = session.query(Document).filter_by(id=doc_id).first()
        
        if not doc:
            return jsonify({
                'success': False,
                'error': '文档不存在'
            }), 404
        
        # 删除相关的文档切片
        slices = session.query(DocumentSlice).filter_by(document_id=doc_id).all()
        for slice_obj in slices:
            session.delete(slice_obj)
        
        session.delete(doc)
        session.commit()
        
        return jsonify({
            'success': True,
            'message': '文档删除成功'
        }), 200
        
    except Exception as e:
        session.rollback()
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500
    finally:
        session.close()

@doc_bp.route('/documents/<doc_id>/slices', methods=['POST'])
def create_document_slice(doc_id):
    """为文档创建新的切片"""
    session = get_postgres_session()
    try:
        # 检查文档是否存在
        doc = session.query(Document).filter_by(id=doc_id).first()
        if not doc:
            return jsonify({
                'success': False,
                'error': '文档不存在'
            }), 404
        
        data = request.get_json()
        
        if not data or not data.get('question') or not data.get('answer'):
            return jsonify({
                'success': False,
                'error': '问题和答案不能为空'
            }), 400
        
        # 生成新的切片ID
        new_id = f"slice{uuid.uuid4().hex[:8]}"
        
        # 获取下一个切片索引
        max_index = session.query(DocumentSlice).filter_by(
            document_id=doc_id
        ).count()
        
        # 创建文档切片
        slice_obj = DocumentSlice(
            id=new_id,
            document_id=doc_id,
            question=data['question'],
            answer=data['answer'],
            slice_index=max_index
        )
        
        session.add(slice_obj)
        session.commit()
        
        return jsonify({
            'success': True,
            'data': {
                'id': slice_obj.id,
                'question': slice_obj.question,
                'answer': slice_obj.answer,
                'sliceIndex': slice_obj.slice_index
            }
        }), 201
        
    except Exception as e:
        session.rollback()
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500
    finally:
        session.close()
